import heapq


class Solution(object):
    def kClosest(self, points, k):
        st = []
        n = len(points)
        for i in range(k):
            st.append([-(points[i][0] ** 2 + points[i][1] ** 2) ** 0.5, points[i][0], points[i][1]])
        heapq.heapify(st)
        for i in range(k, n):
            if (points[i][0] ** 2 + points[i][1] ** 2) ** 0.5 < - st[0][0]:
                heapq.heapreplace(st, [- (points[i][0] ** 2 + points[i][1] ** 2) ** 0.5, points[i][0], points[i][1]])
        return [[t[1], t[2]] for t in st]


data = Solution()
points = [[3, 3], [5, -1], [-2, 4]]
k = 2
print(data.kClosest(points, k))
